
#include <stack>
#include <lexer.h>
#include <ast.h>
#include <parser.h>
#include <unordered_set>
#include <type.h>
#include <iostream>


static int to_integer(const std::string& s) { return atoi(&s[0]); }

static bool check_func_match(FuncDeclaretionStat& decl, FunctionStat& def) {
  return true;
}

Stat Parser::parseContinueStat() {
  _lexer.next_token();
  _lexer.next_token_of_kind(kTokenSemi);
  return ContinueStat{};
}

Stat Parser::parseBreakStat() {
  _lexer.next_token();
  _lexer.next_token_of_kind(kTokenSemi);
  return BreakStat{};
}

Stat Parser::parseBlockStat() {
  if (_lexer.look_ahead() != kTokenLeftCurly) {
    return parseStat();
  }
  BlockStat stat;
  _lexer.next_token();
  _nested++;
  while (_lexer.look_ahead() != kTokenRightCurly) {
    stat.stats.push_back(parseStat());
  }
  _nested--;
  _lexer.next_token_of_kind(kTokenRightCurly);
  return stat;
}

Stat Parser::parseWhileStat() {
  WhileStat stat;
  _lexer.next_token();
  assert(_lexer.look_ahead() == kTokenLeftParen);
  //_lexer.next_token_of_kind(kTokenLeftParen);
  stat.cond = parseExprStat();
  //_lexer.next_token_of_kind(kTokenRightParen);
  stat.stats = parseBlockStat();
  return stat;
}

Stat Parser::parseDoWhileStat() {
  DoWhileStat stat;
  _lexer.next_token();
  stat.stats = parseBlockStat();
  _lexer.next_token_of_kind(kTokenWhile);
  assert(_lexer.look_ahead() == kTokenLeftParen);
  //_lexer.next_token_of_kind(kTokenLeftParen);
  stat.cond = parseExprStat();
  //_lexer.next_token_of_kind(kTokenRightParen);
  return stat;
}

Stat Parser::parseIfStat() {
  IfStat stat;

  _lexer.next_token_of_kind(kTokenIf);

  assert(_lexer.look_ahead() == kTokenLeftParen);
  //_lexer.next_token_of_kind(kTokenLeftParen);
  stat.cond = parseExprStat();
  //_lexer.next_token_of_kind(kTokenRightParen);
  stat.true_stat = parseBlockStat();

  if (_lexer.look_ahead() == kTokenElse) {
    _lexer.next_token();
    stat.false_stat = parseBlockStat();
  }
  return stat;
}

Stat Parser::parseForStat() {
  ForStat stat;
  _lexer.next_token();

  _lexer.next_token_of_kind(kTokenLeftParen);
  stat.init = parseExprStat(kTokenSemi);
  stat.cond = parseExprStat(kTokenSemi);
  stat.step = parseExprStat(kTokenRightParen);
  stat.stat = parseBlockStat();
  return stat;
}

Stat Parser::parseCaseStat() {
  _lexer.next_identifier();
  CaseStat stat;
  Stat exp = parsePrimaryStat();
  assert(exp.type() == typeid(LiteralIntegerStat));
  stat.exp = std::any_cast<LiteralIntegerStat>(exp);
  _lexer.next_token_of_kind(kTokenColon);
  stat.stats = parseSwitchCaseBlockStat();
  return stat;
}

// FIXME
// switch (a) { case A: { abc; }}
Stat Parser::parseSwitchStat() {
  SwitchStat stat;
  _lexer.next_token();
  assert(_lexer.look_ahead() == kTokenLeftParen);
  //_lexer.next_token_of_kind(kTokenLeftParen);
  stat.cond = parseExprStat();
  //_lexer.next_token_of_kind(kTokenRightParen);

  _lexer.next_token_of_kind(kTokenLeftCurly);

  while (_lexer.look_ahead() == kTokenCase) {
    stat.cases.push_back(parseCaseStat());
  }

  if (_lexer.look_ahead() == kTokenDefault) {
    _lexer.next_token();
    _lexer.next_token_of_kind(kTokenColon);
    stat.def = parseSwitchCaseBlockStat();
  }
  _lexer.next_token_of_kind(kTokenRightCurly);
  return stat;
}

Stat Parser::parseSwitchCaseBlockStat() {
  
  if (_lexer.look_ahead() == kTokenLeftCurly) {
    return parseBlockStat();
  }
  BlockStat stat;
  int kind;
  while ((kind = _lexer.look_ahead()) != kTokenCase &&
         kind != kTokenRightCurly && kind != kTokenDefault) {
    stat.stats.push_back(parseStat());
  }
  return stat;
}

Stat Parser::parseStat() {
  bool is_static = false;
  switch (_lexer.look_ahead()) {
    case kTokenLeftCurly: return parseBlockStat();
    case kTokenSemi: return parseEmptyStat();
    case kTokenWhile: return parseWhileStat();
    case kTokenFor: return parseForStat();
    case kTokenIf: return parseIfStat();
    case kTokenDo: return parseDoWhileStat();
    case kTokenContinue: return parseContinueStat();
    case kTokenBreak: return parseBreakStat();
    case kTokenTypedef: return parseTypedefStat();
    case kTokenGoto: return parseGotoStat();
    case kTokenCase: return parseCaseStat();
    case kTokenSwitch: return parseSwitchStat();
    case kTokenReturn: return parseReturnStat();
    case kTokenStatic: is_static = true, _lexer.next_token();
    case kTokenSigned:
    case kTokenUnsigned:
    case kTokenVoid:
    case kTokenChar:
    case kTokenInt: 
    case kTokenLong:
    case kTokenDouble:  {
      std::pair<Type, std::string> p = parseDeclarationSpec();
      //std::string name = _lexer.next_identifier().name;
      if (_lexer.look_ahead() == kTokenLeftParen) return parseFunctionStat(p.first, p.second);
      assert(p.first.type() != typeid(TypeBuiltin) || std::any_cast<TypeBuiltin>(p.first).size() > 0);
      return parseVariableDefineStat(p.first, p.second, is_static || _nested == 0);
    }
    case kTokenStruct: 
      return parseStructStat();
    case kTokenExtern:
    case kTokenEnum: 
    case kTokenFunc: 
    case kTokenVar: 
    case kTokenAuto:return EmptyStat{};
    case kTokenIdentifier: 
    default: return parseExprStat();
  }
  return EmptyStat{};
}


Stat Parser::parseStructStat() {
  StructStat stat;
  _lexer.next_token();
  stat.name = _lexer.next_identifier().name;

  std::vector<std::pair<Type, std::string>> type;
  assert(_lexer.look_ahead() == kTokenLeftCurly);
  _lexer.next_token();
  while (_lexer.look_ahead() != kTokenRightCurly) {
    std::pair<Type, std::string> member = parseDeclarationSpec();
    type.push_back(member);
    stat.member.push_back(member.second);
    _lexer.next_token_of_kind(kTokenSemi);
  }
  _lexer.next_token_of_kind(kTokenRightCurly);
  _lexer.next_token_of_kind(kTokenSemi);

  stat.type = TypeStruct(type);

  return stat;
}

Stat Parser::parseUnionStat() {
  UnionStat stat;
  return stat;
}

// enum [name] { A [= 1, B = 2, C, D]};
Stat Parser::parseEnumStat() {
  EnumStat stat;
  _lexer.next_token();

  if (_lexer.look_ahead() == kTokenIdentifier) {
    stat.name = _lexer.next_token().name;
  }
  int value = 0;
  _lexer.next_token_of_kind(kTokenLeftCurly);

  while (_lexer.look_ahead() != kTokenRightCurly) {
    Token token = _lexer.next_token_of_kind(kTokenIdentifier);
    if (_lexer.look_ahead() == kTokenOpAssign) {
      Stat val = parsePrimaryStat();
      assert(val.type() == typeid(LiteralIntegerStat));
      value = std::any_cast<LiteralIntegerStat>(val).val;
    }
    assert(_enum.count(token.name) == 0);
    _enum[token.name] = value;
    stat.elements[token.name] = value++;

    if (_lexer.look_ahead() == kTokenOpComma) {
      _lexer.next_token();
    }
  }
  _lexer.next_token_of_kind(kTokenRightCurly);
  _lexer.next_token_of_kind(kTokenSemi);
  return stat;
}

// typdef raw alias;
Stat Parser::parseTypedefStat() {
  _lexer.next_token();
  std::string raw, alias, t;
  while (_lexer.look_ahead() != kTokenSemi) {
    t = _lexer.next_token().name;
    if (_lexer.look_ahead() == kTokenSemi) alias = t;
    else {
      if (raw.empty()) raw = t;
      else raw += " " + t;
    }
  }
  _lexer.next_token_of_kind(kTokenSemi);

  assert(_typeAlias.count(alias) == 0);

  _typeAlias[alias] = raw;
  return TypedefStat{raw, alias};
}

// goto label;
Stat Parser::parseGotoStat() {
  _lexer.next_token();
  std::string label = _lexer.next_identifier().name;
  _lexer.next_token_of_kind(kTokenSemi);
  return GotoStat{label};
}

Stat Parser::parseReturnStat() {
  _lexer.next_token();
  return ReturnStat{parseExprStat()};
}


Stat Parser::parseEmptyStat() {
  _lexer.next_token_of_kind(kTokenSemi);
  return EmptyStat{};
}

Stat calcUnaryOpStat(UnaryOpStat stat) {

  if (stat.op == kOpSizeof) {
    return LiteralIntegerStat{0};
  }

  if (stat.exp.type() != typeid(LiteralIntegerStat)) return stat;

  int val = std::any_cast<LiteralIntegerStat>(stat.exp).val;
  switch (stat.op) {
    case kOpBitNot: return LiteralIntegerStat{~val};
    case kOpLogicNot: return LiteralIntegerStat{!val};
    case kOpNegative: return LiteralIntegerStat{-val};
    case kOpPostive: return LiteralIntegerStat{val};
    default: break;
  }
  return stat;
}

Stat calcBinaryOpStat(BinaryOpStat stat) {
  if (stat.lhs.type() != typeid(LiteralIntegerStat)) return stat;
  if (stat.rhs.type() != typeid(LiteralIntegerStat)) return stat;

  int lhs = std::any_cast<LiteralIntegerStat>(stat.lhs).val;
  int rhs = std::any_cast<LiteralIntegerStat>(stat.rhs).val;

  switch (stat.op) {
    case kOpAdd:
      return LiteralIntegerStat{lhs + rhs};
    case kOpSub:
      return LiteralIntegerStat{lhs - rhs};
    case kOpMult:
      return LiteralIntegerStat{lhs * rhs};
    case kOpDiv:
      return LiteralIntegerStat{lhs / rhs};
    case kOpBitOr:
      return LiteralIntegerStat{lhs | rhs};
    case kOpBitAnd:
      return LiteralIntegerStat{lhs & rhs};
    case kOpLogicAnd:
      return LiteralIntegerStat{lhs && rhs};
    case kOpLogicOr:
      return LiteralIntegerStat{lhs || rhs};
    case kOpGreatThan:
      return LiteralIntegerStat{lhs > rhs};
    case kOpGreatEqual:
      return LiteralIntegerStat{lhs >= rhs};
    case kOpLessThan:
      return LiteralIntegerStat{lhs < rhs};
    case kOpLessEqual:
      return LiteralIntegerStat{lhs <= rhs};
    case kOpEqual:
      return LiteralIntegerStat{lhs == rhs};
    case kOpNotEqual:
      return LiteralIntegerStat{lhs != rhs};
    case kOpShiftLeft:
      return LiteralIntegerStat{lhs << rhs};
    case kOpShiftRight:
      return LiteralIntegerStat{lhs >> rhs};
    default:
      break;
  }
  return stat;
}

// 语义分析
// 确定表达式的合法性
// 常量表达式进行替换和计算等
Stat calcConstantStat(Stat stat) {
  if (stat.type() == typeid(UnaryOpStat))  {
    UnaryOpStat unop = std::any_cast<UnaryOpStat>(stat);
    unop.exp = calcConstantStat(unop.exp);
    return calcUnaryOpStat(unop);
  }
  if (stat.type() == typeid(BinaryOpStat)) {
    BinaryOpStat binop = std::any_cast<BinaryOpStat>(stat);
    binop.lhs = calcConstantStat(binop.lhs);
    binop.rhs = calcConstantStat(binop.rhs);
    return calcBinaryOpStat(binop);
  }
  if (stat.type() == typeid(TernaryOpStat)) {
    TernaryOpStat terop = std::any_cast<TernaryOpStat>(stat);
    terop.cond = calcConstantStat(terop.cond);
    terop.true_stat = calcConstantStat(terop.true_stat);
    terop.false_stat = calcConstantStat(terop.false_stat);
    if (terop.cond.type() == typeid(LiteralIntegerStat)) {
      int val = std::any_cast<LiteralIntegerStat>(terop.cond).val;
      return val ? terop.true_stat : terop.false_stat;
    }
    return terop;
  }
  return stat;
}

// 语法分析
// 根据计算规则生成抽象语法树
void calcExpr(std::stack<Stat>& stats, std::stack<Operator>& ops, int priority) {
  while (!ops.empty()) {
    Operator sp = ops.top();
    int pri = operatorPriority(sp);
    OpAssociative asso = operatorAssociative(sp);
    OpType type = operatorType(sp);
    if (priority < pri || (priority == pri && asso == kRightToLeft)) {
      break;
    }
    ops.pop();
    if (type == kUnaryOp) {
      stats.top() = UnaryOpStat{sp, std::move(stats.top())};
    } else if (type == kBinaryOp) {
      BinaryOpStat bstat;
      bstat.op = sp;
      bstat.rhs = std::move(stats.top());
      stats.pop();
      bstat.lhs = std::move(stats.top());
      stats.top() = std::move(bstat);
    } else if (type == kTernaryOp) {
      TernaryOpStat tstat;
      tstat.false_stat = std::move(stats.top());
      stats.pop();
      tstat.true_stat = std::move(stats.top());
      stats.pop();
      tstat.cond = std::move(stats.top());
      stats.top() = std::move(tstat);
    } else {
      abort();
    }
  }
  // TODO
}

Stat Parser::parseExprStat(TokenType terminal, TokenType ter2) {
  std::stack<Stat> stats;
  std::stack<Operator> ops;
  if (_lexer.look_ahead() == kTokenLeftParen)
    terminal = kTokenRightParen, _lexer.next_token();
  else if (_lexer.look_ahead() == kTokenLeftBrack)
    terminal = kTokenRightBrack, _lexer.next_token();

  int kind;
  while ((kind = _lexer.look_ahead()) != terminal && kind != ter2 && kind != kTokenEof) {
    Operator op = kOpNone;
    switch (_lexer.look_ahead()) {
      case kTokenColon: 
        break;
      case kTokenMemberAccess:
        _lexer.next_token();
        stats.top() = BinaryOpStat{kOpMemberAccess, std::move(stats.top()), NameStat{_lexer.next_identifier().name}};
        break;
      case kTokenOpPointer:
        _lexer.next_token();
        stats.top() = BinaryOpStat{kOpPointer, std::move(stats.top()), NameStat{_lexer.next_identifier().name}};
        break;
      case kTokenLeftBrack:
        stats.top() = BinaryOpStat{kOpArraySubscript, std::move(stats.top()), parseExprStat(kTokenRightBrack)};
        break;

      case kTokenLiteraltInteger:
        stats.push(LiteralIntegerStat{to_integer(_lexer.next_token().name)});
        break;
      case kTokenIdentifier: {
        std::string name = _lexer.next_token().name;
        if (_lexer.look_ahead() == kTokenColon) {
          _lexer.next_token();
          return LabelStat{name};
        }
        if (_enum.count(name)) stats.push(LiteralIntegerStat{_enum[name]});
        else if (_vars.count(name)) stats.push(VarStat{name});
        else if (_func.count(name)) {
          FuncCallArgsStat args;
          _lexer.next_token_of_kind(kTokenLeftParen);
          while (_lexer.look_back().kind != kTokenRightParen) {
            args.args.push_back(parseExprStat(kTokenRightParen, kTokenOpComma));
          }
          stats.push(BinaryOpStat{kOpFuncCall, NameStat{name}, args, std::any_cast<FunctionStat>(_func[name]).ret});
        } else panic("undefined identifier!");
        break;
      }
        break;
      case kTokenLeftParen:
        stats.push(parseExprStat());
        break;
      case kTokenSelfDec:
        _lexer.next_token();
        if (!stats.empty() && statWritable(stats.top())) {
          stats.top() = UnaryOpStat{kOpSelfDecSufix, std::move(stats.top())};
        } else {
          op = kOpSelfDecPrefix;
        }
        break;
      case kTokenSelfInc:
        _lexer.next_token();
        if (!stats.empty() && statWritable(stats.top())) {
          stats.top() = UnaryOpStat{kOpSelfIncSufix, std::move(stats.top())};
        } else {
          op = kOpSelfIncPrefix;
        }
        break;
      case kTokenStar: // mult or ptrderef
        _lexer.next_token();
        op = (stats.size() == ops.size()) ? kOpPtrDeref : kOpMult;
        break;
      case kTokenAmp: // bitand or address
        _lexer.next_token();
        op = (stats.size() == ops.size()) ? kOpAddress : kOpBitAnd;
        break;
      case kTokenMinus:
        _lexer.next_token();
        op = (stats.size() == ops.size()) ? kOpNegative : kOpSub;
        break;
      case kTokenPlus:
        _lexer.next_token();
        op = (stats.size() == ops.size()) ? kOpPostive : kOpAdd;
        break;
      case kTokenQuestion: 
      case kTokenOpLogicNot:
      default:
        op = genOperator(_lexer.next_token().kind);
        break;
    }
    if (op != kOpNone) {
      calcExpr(stats, ops, operatorPriority(op));
      if (op == kOpCond) {
        stats.push(parseExprStat(kTokenColon));
      }
      ops.push(op);
    }
  }
  calcExpr(stats, ops, kLowestPriority);
  _lexer.next_token();
  return calcConstantStat(stats.top());
}

Stat Parser::parsePrimaryStat() {
  if (_lexer.look_ahead() == kTokenLeftParen) {
    _lexer.next_token();
    Stat stat = parseExprStat();
    _lexer.next_token_of_kind(kTokenRightParen);
    return stat;
  } 
  if (_lexer.look_ahead() == kTokenLiteraltInteger) {
    return LiteralIntegerStat{to_integer(_lexer.next_token().name)};
  }
  if (_lexer.look_ahead() == kTokenIdentifier) {
    return NameStat{_lexer.next_token().name};
  }

  return Stat{};
}

bool Parser::statWritable(Stat stat) {
  if (stat.type() == typeid(VarStat)) {
      const std::string& name = std::any_cast<VarStat>(stat).name;
      return _vars.count(name);
  }

  if (stat.type() == typeid(BinaryOpStat)) {
    Operator op = std::any_cast<BinaryOpStat>(stat).op;
    return op == kOpArraySubscript || op == kOpMemberAccess || op == kOpPointer;
  }
  return false;
}

Stat Parser::parse() {
  BlockStat stat;
  while (_lexer.look_ahead() != kTokenEof) {
    stat.stats.push_back(parseStat());
  }
  return stat;
}

std::pair<Type, std::string> Parser::parseDeclarationSpec() {
  Type type;
  switch (_lexer.look_ahead()) {
    case kTokenVoid: type = TypeBuiltin::Void; break;
    case kTokenBool: type = TypeBuiltin::Bool; break;
    case kTokenChar: type = TypeBuiltin::Char; break;
    case kTokenShort: type = TypeBuiltin::Short; break;
    case kTokenInt: type = TypeBuiltin::Int; break;
    case kTokenLong: type = TypeBuiltin::Long; break;
    case kTokenFloat: type = TypeBuiltin::Float; break;
    case kTokenDouble: type = TypeBuiltin::Double; break;
    default:  {
    }
  }
  _lexer.next_token();
  while (_lexer.look_ahead() == kTokenStar) {
    type = TypePointer{std::move(type)};
    _lexer.next_token();
  }
  std::string identifier = _lexer.next_identifier().name;

  while (_lexer.look_ahead() == kTokenLeftBrack) {
    int dlen = std::any_cast<LiteralIntegerStat>(parseExprStat()).val;
    type = TypeArray{std::move(type), (size_t)dlen};
  }
  return {type, identifier};
}

Stat Parser::parseFunctionStat(Type ret, std::string name) {
  FunctionStat stat;
  _locals.clear();
  stat.ret = ret;

  _lexer.next_token_of_kind(kTokenLeftParen);
  while (_lexer.look_ahead() != kTokenRightParen) {
    std::pair<Type, std::string> p = parseDeclarationSpec();
    stat.args.push_back(p);
    _vars.insert(p.second);
    if (_lexer.look_ahead() == kTokenOpComma) _lexer.next_token();
  }
  _lexer.next_token_of_kind(kTokenRightParen);
  assert(_lexer.look_ahead() == kTokenLeftCurly);
  stat.name = name;
  stat.stats = std::move(parseBlockStat());
  stat.local = std::move(_locals);
  _func[name] = stat;
  return stat;
}

Stat Parser::parseVariableDefineStat(Type type, std::string name, bool is_static) {
  Type base = type;
  while (base.type() == typeid(TypePointer)) {
    base = std::any_cast<TypePointer>(base)._base;
  }
  VariableDefineStat stat;

  _lexer.next_token();
  if (_lexer.look_back().kind == kTokenOpAssign) {
    stat.stat.push_back(VariableStat{type, name, parseExprStat(kTokenOpComma), is_static});
  } else {
    stat.stat.push_back(VariableStat{type, name, Stat{}, is_static});
  }
  _locals.push_back({type, name});
  _vars.insert(name);

  while (_lexer.look_back().kind != kTokenSemi) {
    type = base;
    while (_lexer.look_ahead() == kTokenStar) {
      type = TypePointer{std::move(type)};
      _lexer.next_token();
    }
    name = _lexer.next_identifier().name;

    while (_lexer.look_ahead() == kTokenLeftBrack) {
      int dlen = std::any_cast<LiteralIntegerStat>(parseExprStat()).val;
      type = TypeArray{std::move(type), (size_t)dlen};
    }
    _lexer.next_token();
    if (_lexer.look_back().kind == kTokenOpAssign) {
      stat.stat.push_back(VariableStat{type, name, parseExprStat(kTokenOpComma), is_static});
    } else {
      stat.stat.push_back(VariableStat{type, name, Stat{}, is_static});
    }
    _locals.push_back({type, name});
    _vars.insert(name);
  }

  return stat;
}
